package com.yupi.loveplanneraiagent.rag;

import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

import java.util.List;

import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

// 为方便开发调试和部署，临时注释，如果需要使用 PgVector 存储知识库，取消注释即可
/**
 * pgsql建表语句
 * -- 创建表
 * CREATE TABLE vector_store (
 *     id UUID default uuid_generate_v4() not null primary key,  -- UUID主键，默认生成随机UUID
 *     content TEXT,                                   -- 文本内容字段
 *     metadata JSON,                                  -- JSON元数据字段
 *     embedding vector(1536)                               -- 向量嵌入字段
 * );
 */
//@Configuration
@Slf4j
public class PgVectorVectorStoreConfig {

    @Resource
    private LoveAppDocumentLoader loveAppDocumentLoader;

    @Bean
    public VectorStore pgVectorVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel dashscopeEmbeddingModel) {
        VectorStore vectorStore = PgVectorStore.builder(jdbcTemplate, dashscopeEmbeddingModel)
                .dimensions(1536)                    // Optional: defaults to model dimensions or 1536
                .distanceType(COSINE_DISTANCE)       // Optional: defaults to COSINE_DISTANCE
                .indexType(HNSW)                     // Optional: defaults to HNSW
                .initializeSchema(true)              // Optional: defaults to false
                .schemaName("public")                // Optional: defaults to "public"
                .vectorTableName("vector_store")     // Optional: defaults to "vector_store"
                .maxDocumentBatchSize(10000)         // Optional: defaults to 10000
                .build();
        // 检查是否已有文档，如果没有则加载文档
        loadDocumentsIfEmpty(vectorStore);
        return vectorStore;
    }

    /**
     * 检查向量存储中是否已有文档，如果没有则加载文档
     * @param vectorStore 向量存储实例
     */
    private void loadDocumentsIfEmpty(VectorStore vectorStore) {
        try {
            // 尝试执行一个简单的查询来检查是否有文档
            // 使用一个简单的相似性搜索来检查是否已有文档
            List<Document> testSearch = vectorStore.similaritySearch("我已经结婚了，但是婚后关系不太亲密，怎么办？");
            // 如果没有找到文档，说明数据库为空，需要加载文档
            if (testSearch.isEmpty()) {
                log.info("向量存储中未发现文档，开始加载文档...");
                List<Document> documents = loveAppDocumentLoader.loadMarkdowns();
                if (documents.isEmpty()) {
                    log.warn("文档加载器返回空文档列表，跳过加载");
                    return;
                }
                log.info("检测到 {} 个文档，开始加载到向量存储", documents.size());
                // 如果文档数量超过限制，则分批处理
                if (documents.size() > 25) {
                    log.info("文档数量超过单批处理限制，将分批处理以避免输入限制");
                    // 分批添加文档，每批最多25个文档
                    for (int i = 0; i < documents.size(); i += 25) {
                        int endIndex = Math.min(i + 25, documents.size());
                        List<Document> batch = documents.subList(i, endIndex);
                        vectorStore.add(batch);
                        log.info("已处理文档批次: {}/{}",
                                (i / 25 + 1),
                                (int) Math.ceil((double) documents.size() / 25));
                    }
                } else {
                    vectorStore.add(documents);
                    log.info("所有文档已加载完成");
                }
                log.info("文档加载完成，共加载 {} 个文档", documents.size());
            } else {
                log.info("向量存储中已存在文档，跳过加载步骤");
            }
        } catch (Exception e) {
            log.error("检查或加载文档时出错: {}", e.getMessage(), e);
            // 如果出现错误，可以选择重新尝试加载或抛出异常
            // 这里我们记录错误但不中断应用启动
        }
    }
}
